import { BaseAgentBackend, BackendTypes } from "../base";
import { getTextEmbedding as _getTextEmbedding } from "./embedding";
import { getChatCompletion as _getChatCompletion } from "./caht";
import { modelTokenLimit } from "./config";

export class BaiDuBackend extends BaseAgentBackend {
    static async getTextEmbedding(
        request: BackendTypes.TextEmbeddingRequest
    ): Promise<BackendTypes.TextEmbeddingResponse> {
        const response = await _getTextEmbedding(
            { ...request, input: [request.input] },
            request.apiKey
        );
        return { ...response, model: request.model };
    }
    static async getChatCompletion(
        request: BackendTypes.ChatCompletionRequest,
        modelConfig: Record<string, any>
    ): Promise<BackendTypes.ChatCompletionResponse> {
        const system = request.messages[0].content;
        const messages = request.messages.slice(1).map((message) => {
            if (message.role === "system") {
                throw new Error("System message is not allowed in chat completion");
            }
            return message;
        }) as any;
        const request2 = {
            model: request.model,
            messages,
            system
        };
        const response = await _getChatCompletion(request2, request.apiKey, modelConfig);
        return {
            id: response.id,
            created: response.created,
            model: request.model,
            message: { role: "assistant", content: response.result }
        };
    }
    static getTokenLimit(modelType: string): number {
        const result = modelTokenLimit[modelType];
        if (result) {
            return result;
        } else {
            throw new Error("Invalid model type");
        }
    }
}
